package com.bnd.service.impl;

import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.Role;
import com.bnd.config.AiConfig;
import com.bnd.config.SseEmitterUTF8;
import com.bnd.domain.session.ChatSession;
import com.bnd.service.ChatHistoryService;
import com.bnd.service.RoleChatService;
import com.bnd.service.StreamChatHandle;
import io.reactivex.Flowable;
import io.reactivex.schedulers.Schedulers;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.util.List;
import java.util.concurrent.CountDownLatch;

@Slf4j
@Service
public class RoleChatServiceImpl implements RoleChatService {
    @Resource
    private AiConfig aiConfig;
    @Resource
    private ChatHistoryService chatHistoryService;
    private static final String CHARACTER_MODEL = "qwen-plus-character";

    @Override
    public SseEmitter doStreamChat(String userInput, ChatSession chatSession) {
        //1. 初始化参数
        SseEmitterUTF8 emitter = chatSession.getSseEmitter();
        Generation gen = chatSession.getGeneration();
        CountDownLatch latch = new CountDownLatch(1);
        List<Message> messageList = chatSession.getMessageList();

        // 将用户输入加入消息列表
        messageList.add(Message.builder()
                .role(Role.USER.getValue())
                .content(userInput)
                .build());

        GenerationParam param = GenerationParam.builder()
                .apiKey(aiConfig.getDashscope().getApiKey())
                .model(CHARACTER_MODEL)
                .messages(messageList)
                .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                .incrementalOutput(true) // 开启增量输出，流式返回
                .build();

        //2. 打字机推送
        try {
            Flowable<GenerationResult> result = gen.streamCall(param);
            StringBuilder fullContent = new StringBuilder();
            result
                    .subscribeOn(Schedulers.io()) // IO线程执行请求
                    .observeOn(Schedulers.computation()) // 计算线程处理响应
                    .subscribe(
                            // onNext: 处理每个响应片段
                            message -> {
                                try {
                                    String content = message.getOutput().getChoices().get(0).getMessage().getContent();
                                    String finishReason = message.getOutput().getChoices().get(0).getFinishReason();

                                    // 输出内容到SSE
                                    emitter.send(content);
                                    fullContent.append(content);

                                    // 当 finishReason 不为 null 时，表示是最后一个 chunk，输出用量信息
                                    if (finishReason != null && !"null".equals(finishReason)) {
                                        log.info("--- 请求用量 ---");
                                        log.info("输入 Tokens：{}", message.getUsage().getInputTokens());
                                        log.info("输出 Tokens：{}", message.getUsage().getOutputTokens());
                                        log.info("总 Tokens：{}", message.getUsage().getTotalTokens());

                                        // 可以选择发送用量信息给前端
                                        // emitter.send(SseEmitter.event().name("usage").data(message.getUsage()));
                                    }
                                } catch (Exception e) {
                                    log.error("处理消息片段异常", e);
                                    emitter.completeWithError(e);
                                }
                            },
                            // onError: 处理错误
                            error -> {
                                log.error("请求失败: {}", error.getMessage(), error);
                                try {
                                    emitter.completeWithError(error);
                                } catch (Exception e) {
                                    log.error("关闭SSE连接异常", e);
                                }
                                latch.countDown();
                            },
                            // onComplete: 完成回调
                            () -> {
                                log.info("完整响应: {}", fullContent.toString());
                                messageList.add(Message.builder()
                                        .role(Role.ASSISTANT.getValue())
                                        .content(fullContent.toString())
                                        .build());
                                try {
                                    emitter.complete();
                                } catch (Exception e) {
                                    log.error("完成SSE连接异常", e);
                                }
                                latch.countDown();
                            }
                    );
            // 主线程等待异步任务完成
            latch.await();
        } catch (Exception e) {
            log.error("流式推送异常: {}", e.getMessage(), e);
            try {
                emitter.completeWithError(e);
            } catch (Exception ex) {
                log.error("关闭SSE连接异常", ex);
            }
        }
        return emitter;
    }

    @Override
    public StreamChatHandle doStreamChatDeferComplete(String userInput, ChatSession chatSession) {
        SseEmitterUTF8 emitter = chatSession.getSseEmitter();
        Generation gen = chatSession.getGeneration();
        List<Message> messageList = chatSession.getMessageList();

        messageList.add(Message.builder().role(Role.USER.getValue()).content(userInput).build());

        GenerationParam param = GenerationParam.builder()
                .apiKey(aiConfig.getDashscope().getApiKey())
                .model(CHARACTER_MODEL)
                .messages(messageList)
                .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                .incrementalOutput(true)
                .build();

        java.util.concurrent.CompletableFuture<String> textDone = new java.util.concurrent.CompletableFuture<>();

        try {
            Flowable<GenerationResult> result = gen.streamCall(param);
            StringBuilder fullContent = new StringBuilder();
            result
                    .subscribeOn(Schedulers.io())
                    .observeOn(Schedulers.computation())
                    .subscribe(
                            msg -> {
                                try {
                                    String content = msg.getOutput().getChoices().get(0).getMessage().getContent();
                                    emitter.send(SseEmitter.event().name("text").data(content));
                                    fullContent.append(content);
                                } catch (Exception e) {
                                    log.error("处理消息片段异常", e);
                                    emitter.completeWithError(e);
                                    textDone.completeExceptionally(e);
                                }
                            },
                            err -> {
                                log.error("请求失败: {}", err.getMessage(), err);
                                try { emitter.completeWithError(err); } catch (Exception ignore) {}
                                textDone.completeExceptionally(err);
                            },
                            () -> {
                                String text = fullContent.toString();
                                log.info("完整响应: {}", text);
                                messageList.add(Message.builder().role(Role.ASSISTANT.getValue()).content(text).build());
                                // 不在这里 complete，由调用方在附加逻辑（如文生图）完成后再关闭
                                textDone.complete(text);
                            }
                    );
        } catch (Exception e) {
            log.error("流式推送异常: {}", e.getMessage(), e);
            try { emitter.completeWithError(e); } catch (Exception ignore) {}
            textDone.completeExceptionally(e);
        }

        return new StreamChatHandle(emitter, textDone);
    }

    @Override
    public void closeChatSession(ChatSession removedSession) {
        //1. 关闭资源
        if (removedSession != null) {
            // 关闭SSE连接（如果还存在）
            try {
                SseEmitter emitter = removedSession.getSseEmitter();
                if (emitter != null) {
                    emitter.complete();
                }
                //2. 保存聊天记录
                chatHistoryService.batchSaveChatHistory(removedSession.getChatId(), removedSession.getCharacterId(), removedSession.getMessageList());
            } catch (Exception e) {
                log.warn("关闭SSE连接时出现异常: {}", e.getMessage());
            }
        }
    }
}
